import argparse
from tools.draw_data import drawdata
from tools.Detection import Detection
from mindspore.train.serialization import load


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='OM', help='choose model')
    parser.add_argument('--output', type=str, default='output/', help='output path')
    parser.add_argument('--conf_thres', type=float, default=0.6, help='object confidence threshold')
    parser.add_argument('--iou_thres', type=float, default=0.5, help='IOU threshold for NMS')
    parser.add_argument('--target_weights', type=str, default='model_data/ssd_target_512_sim_t.om',
                        help='target detection weights path')
    parser.add_argument('--helmet_weights', type=str, default='head/model_data/ssd_helmet_512_sim_t.om',
                        help='helmet detection weights path')
    parser.add_argument('--pipeline_target_path', type=str, default='pipeline/ssd_target_512.pipeline',
                        help='helmet pipeline path')
    parser.add_argument('--pipeline_helmet_path', type=str, default='pipeline/ssd_helmet_512.pipeline',
                        help='helmet pipeline path')
    parser.add_argument('--input_shape', type=int, default=512, help='target model input shape')
    parser.add_argument('--input_shape2', type=int, default=512, help='helmet model input shape')

    parser.add_argument('--predict', action='store_true', default=False, help='predict')
    parser.add_argument('--draw_data', action='store_true', default=False, help='Draw the data')
    parser.add_argument('--video', action='store_true', default=False, help='predict video')
    parser.add_argument('--video_path', type=str, default='input/video/test.avi', help='video_path')
    parser.add_argument('--image', action='store_true', default=False, help='predict image')

    opt = parser.parse_args()
    print(opt)

    if opt.model:
        model_weights = load(opt.target_weights)
        model_helmet = load(opt.target_helmet)

    if opt.predict:
        Detection(opt)

    if opt.draw_data:
        drawdata(opt)
